from camel.models import ModelFactory
from camel.societies import RolePlaying
from camel.types import ModelPlatformType
from dotenv import load_dotenv
import os
load_dotenv()

Model_Type=os.getenv("MODEL_TYPE")
Model_Api=os.getenv("ZHIPU_API_KEY")
Model_Url=os.getenv("MODEL_URL")
model = ModelFactory.create(
    model_platform=ModelPlatformType.OPENAI,
    model_type=Model_Type,
    api_key=Model_Api,
    url=Model_Url
)
task_kwargs = {
    'task_prompt':'制定一个新人小白学习camel的教程，请用中文回答',
    'with_task_specify':True,
    'task_specify_agent_kwargs':{
        'model': model,
    }
}
user_role_kwargs = {
    'user_role_name':'一个camel小白',
    'user_agent_kwargs':{
        'model': model,
    }
}
assistant_role_kwargs = {
    'assistant_role_name':'一个camel专家,了解最新的camel功能，接口。',
    'assistant_agent_kwargs':{
        'model': model,
    }
}
society = RolePlaying(
    **task_kwargs,
    **user_role_kwargs,
    **assistant_role_kwargs
)
def is_terminated(response):
    """
    Give alerts when the session should be terminated.
    """
    if response.terminated:
        role = response.msg.role_type.name
        reason = response.info['termination_reasons']
        print(f'AI {role} terminated due to {reason}')

    return response.terminated

def run(society, round_limit: int=10):

    # Get the initial message from the ai assistant to the ai user
    input_msg = society.init_chat()

    # Starting the interactive session
    for _ in range(round_limit):

        # Get the both responses for this round
        assistant_response, user_response = society.step(input_msg)

        # Check the termination condition
        if is_terminated(assistant_response) or is_terminated(user_response):
            break

        # Get the results
        print(f'[AI User] {user_response.msg.content}.\n')

        # 写入一个md文件
        with open('output.md', 'a', encoding='utf-8') as f:
            f.write(f'[AI User] {user_response.msg.content}.\n')

        # Check if the task is end
        if 'CAMEL_TASK_DONE' in user_response.msg.content:
            break
        print(f'[AI Assistant] {assistant_response.msg.content}.\n')

          # 写入一个md文件
        with open('output.md', 'a', encoding='utf-8') as f:
            f.write(f'[AI Assistant] {assistant_response.msg.content}.\n')

        # Get the input message for the next round
        input_msg = assistant_response.msg

    return None
if __name__ == "__main__":
    run(society, round_limit=10)